package backend.opt;

import backend.Instruction.*;
import backend.Register.RISCVreg;
import backend.Register.physicalFReg;
import backend.Register.physicalReg;
import backend.component.RISCVblock;
import backend.component.RISCVfunction;
import backend.component.RISCVmodule;

import java.util.*;
import java.util.stream.Collectors;

/**
 * RISCVScheduler 是一个用于 RISC-V 指令调度的优化器，它通过对指令进行重排序来隐藏指令延迟，提高指令级并行性（ILP），从而优化代码执行效率。
 */
public class RISCVScheduler {
    RISCVmodule riscVmodule;

    // 指令延迟配置
    private static final Map<String, Integer> INSTRUCTION_LATENCY = new HashMap<>() {{
        put("mul", 3);      put("div", 8);      put("rem", 10);     put("divw", 8);
        put("fdiv.s", 12);  put("fmul.s", 4);   put("fsub.s", 4);   put("fadd.s", 4);
        put("fsub.d", 12);  put("fadd.d", 12);  put("fmv.s", 4);    put("fmv.d", 12);
        put("lw", 3);       put("flw", 3);      put("ld", 4);       put("fld", 5);
        put("sw", 3);       put("fsw", 4);      put("sd", 4);       put("fsd", 5);
        put("fcvt.w.s", 5); put("fcvt.s.w", 4); put("fcvt.s.d", 12); put("fcvt.d.s", 12);
        put("sltu", 1);     put("slti", 1);     put("xor", 1);      put("sgt", 1);
        put("slt", 1);      put("feq.s", 12);   put("flt.s", 12);   put("fle.s", 12);
        put("fgt.s", 12);   put("fge.s", 12);   put("add", 1);      put("addi", 1);
        put("fadd", 4);     put("addw", 1);     put("mulw", 3);     put("sub", 1);
        put("fsub", 4);     put("fdiv", 12);    put("fdiv.d", 12);  put("nop", 1);
        put("jal", 3);      put("jalr", 3);     put("beq", 1);      put("bne", 1);
        put("blt", 1);      put("bge", 1);      put("lb", 3);       put("lbu", 3);
        put("sb", 3);       put("li", 1);       put("j", 3);        put("move", 1);
        put("la", 3);       put("jr", 3);       put("sra", 1);      put("subw", 1);
        put("xori", 1);     put("beqz", 1);     put("addiw", 1);    put("sll", 1);
        put("srl", 1);
    }};

    public RISCVScheduler(RISCVmodule riscVmodule) {
        this.riscVmodule = riscVmodule;
    }

    public void run() {
        for (RISCVfunction riscVfunction : riscVmodule.getFunctions()) {
            for (RISCVblock riscVblock : riscVfunction.getBlocks()) {
                if(!hasCall(riscVblock)) {
                    scheduleFunctionBlocks(riscVblock);
                }
            }
        }
    }

    private boolean hasCall(RISCVblock riscVblock){
        for(RISCVinst riscVinst:riscVblock.getInstructionlist()){
            if(riscVinst instanceof RISCVcall){
                return true;
            }
        }
        return false;
    }

    public void scheduleFunctionBlocks(RISCVblock block) {
        LinkedList<RISCVinst> instructions = new LinkedList<>(block.getInstructionlist());

        // 保留最后一条指令（通常是跳转指令）
        if (instructions.isEmpty()) return;

        RISCVinst lastInst = instructions.removeLast();

        // 只对有多条指令的块进行调度
        if (instructions.size() > 1) {
            //instructions.forEach(inst -> System.out.print(inst));
            LinkedList<RISCVinst> optimized = scheduleBasicBlock(instructions);
            optimized.addLast(lastInst);

            // 更新基本块的指令列表
            block.setInstructionlist(optimized);
        } else {
            instructions.addLast(lastInst);
        }
    }

    private LinkedList<RISCVinst> scheduleBasicBlock(LinkedList<RISCVinst> instructions) {
        if (instructions.size() <= 1) return instructions;

        // 构建依赖图
        List<InstructionNode> nodes = buildDependencyGraph(instructions);

        // 计算关键路径
        calculateCriticalPaths(nodes);

        // 执行列表调度
        return listScheduling(nodes);
    }

    private static class InstructionNode {
        RISCVinst inst;
        int index;
        int latency;
        int criticalPath;
        Set<InstructionNode> predecessors = new HashSet<>();
        Set<InstructionNode> successors = new HashSet<>();
        int inDegree = 0;
        int earliestStart = 0;

        public InstructionNode(RISCVinst inst, int index) {
            this.inst = inst;
            this.index = index;
            this.latency = getLatency(inst);
        }
    }

    private List<InstructionNode> buildDependencyGraph(LinkedList<RISCVinst> instructions) {
        List<InstructionNode> nodes = new ArrayList<>();
        Map<RISCVreg, InstructionNode> lastWriter = new HashMap<>();
        Map<RISCVreg, List<InstructionNode>> lastReaders = new HashMap<>();
        List<InstructionNode> memoryOps = new ArrayList<>();

        // 创建节点
        for (int i = 0; i < instructions.size(); i++) {
            nodes.add(new InstructionNode(instructions.get(i), i));
        }

        // 建立依赖关系
        for (InstructionNode node : nodes) {
            RISCVinst inst = node.inst;

            // 处理读取的寄存器
            if (inst.getRegUse() != null) {
                for (RISCVreg useReg : inst.getRegUse()) {
                    if (useReg != null) {
                        // RAW依赖：读取之前必须等待写入完成
                        if (lastWriter.containsKey(useReg)) {
                            addDependency(lastWriter.get(useReg), node);
                        }

                        // 记录读取者
                        lastReaders.computeIfAbsent(useReg, k -> new ArrayList<>()).add(node);
                    }
                }
            }

            // 处理写入的寄存器
            if (inst.getReg() != null) {
                RISCVreg defReg = inst.getReg();

                // WAW依赖：写入之前必须等待前一个写入完成
                if (lastWriter.containsKey(defReg)) {
                    addDependency(lastWriter.get(defReg), node);
                }

                // WAR依赖：写入之前必须等待所有读取完成
                if (lastReaders.containsKey(defReg)) {
                    for (InstructionNode reader : lastReaders.get(defReg)) {
                        if (!reader.equals(node)) {
                            addDependency(reader, node);
                        }
                    }
                }

                // 更新写入者和清空读取者
                lastWriter.put(defReg, node);
                lastReaders.remove(defReg);
            }

            // 在 buildDependencyGraph 方法中修改内存处理部分
            if (isMemoryAccess(inst)) {
                // 只依赖最近的内存操作
                if (!memoryOps.isEmpty()) {
                    InstructionNode lastMemOp = memoryOps.get(memoryOps.size() - 1);
                    addDependency(lastMemOp, node);
                }
                memoryOps.add(node);
            }
        }

        return nodes;
    }

    private void addDependency(InstructionNode from, InstructionNode to) {
        if (from != to && !from.successors.contains(to)) {
            from.successors.add(to);
            to.predecessors.add(from);
            to.inDegree++;
        }
    }

    private void calculateCriticalPaths(List<InstructionNode> nodes) {
        // 正向拓扑排序计算最早开始时间
        Queue<InstructionNode> queue = new LinkedList<>();
        int[] indegreeCopy = new int[nodes.size()];
        for (int i = 0; i < nodes.size(); i++) {
            indegreeCopy[i] = nodes.get(i).inDegree;
            if (nodes.get(i).inDegree == 0) {
                queue.offer(nodes.get(i));
                nodes.get(i).earliestStart = 0;
            }
        }

        // 正向传播
        while (!queue.isEmpty()) {
            InstructionNode current = queue.poll();
            for (InstructionNode succ : current.successors) {
                succ.earliestStart = Math.max(succ.earliestStart, current.earliestStart + current.latency);
                if (--indegreeCopy[succ.index] == 0) {
                    queue.offer(succ);
                }
            }
        }

        // 反向计算关键路径（迭代代替递归）
        int[] criticalPaths = new int[nodes.size()];
        List<InstructionNode> reverseOrder = new ArrayList<>(nodes);
        Collections.reverse(reverseOrder);

        for (InstructionNode node : reverseOrder) {
            int maxSuccPath = 0;
            for (InstructionNode succ : node.successors) {
                maxSuccPath = Math.max(maxSuccPath, criticalPaths[succ.index]);
            }
            criticalPaths[node.index] = node.latency + maxSuccPath;
            node.criticalPath = criticalPaths[node.index];
        }
    }

    private int calculateNodeCriticalPath(InstructionNode node) {
        if (node.successors.isEmpty()) {
            return node.latency;
        }

        int maxSuccessorPath = 0;
        for (InstructionNode successor : node.successors) {
            maxSuccessorPath = Math.max(maxSuccessorPath, calculateNodeCriticalPath(successor));
        }

        return node.latency + maxSuccessorPath;
    }

    private LinkedList<RISCVinst> listScheduling(List<InstructionNode> nodes) {
        LinkedList<RISCVinst> scheduled = new LinkedList<>();
        PriorityQueue<InstructionNode> readyQueue = new PriorityQueue<>((a, b) -> {
            // 首先按关键路径排序（降序）
            int pathCompare = Integer.compare(b.criticalPath, a.criticalPath);
            if (pathCompare != 0) return pathCompare;

            // 如果关键路径相同，按原始顺序排序
            return Integer.compare(a.index, b.index);
        });

        // 重新初始化入度
        for (InstructionNode node : nodes) {
            node.inDegree = node.predecessors.size();
        }

        // 初始化就绪队列
        for (InstructionNode node : nodes) {
            if (node.inDegree == 0) {
                readyQueue.offer(node);
            }
        }

        // 调度指令
        int scheduledCount = 0;
        while (!readyQueue.isEmpty()) {
            InstructionNode selected = readyQueue.poll();
            scheduled.add(selected.inst);
            scheduledCount++;

            // 更新后继节点
            for (InstructionNode successor : selected.successors) {
                if (--successor.inDegree == 0) {
                    readyQueue.offer(successor);
                }
            }

            // 防止死循环：当调度数量超过节点数时中断
            if (scheduledCount > nodes.size()) {
                System.err.println("Deadlock detected in scheduling. Returning original order.");
                return nodes.stream()
                        .map(node -> node.inst)
                        .collect(Collectors.toCollection(LinkedList::new));
            }
        }

        // 检查是否全部调度
        if (scheduled.size() != nodes.size()) {
            System.err.println("Incomplete scheduling. Returning original order.");
            return nodes.stream()
                    .map(node -> node.inst)
                    .collect(Collectors.toCollection(LinkedList::new));
        }

        return scheduled;
    }

    private static int getLatency(RISCVinst inst) {
        String type = inst.toString().trim();
        int spaceIndex = type.indexOf(' ');
        if (spaceIndex > 0) {
            type = type.substring(0, spaceIndex);
        }

        return INSTRUCTION_LATENCY.getOrDefault(type, 1);
    }

    private boolean isMemoryAccess(RISCVinst inst) {
        // 包含 load/store 和函数调用
        if (inst instanceof RISCVload || inst instanceof RISCVstore) {
            return true;
        }
        // 识别函数调用指令
        String instStr = inst.toString().trim();
        return instStr.startsWith("jal") || instStr.startsWith("call");
    }
}

